/**
 * @fileoverview Typegraph visualizer for pytype.
 *
 * This file is intended to be injected into the Jinja2 template for the
 * visualizer, in order to keep that as a single file. It's a separate file to
 * take advantage of JS tools and IDE integration.
 *
 * IDs are very important in this. There are two kinds:
 * - Serialized ID, the ID number for a part of the pytype typegraph.
 * - Cytoscape ID, the value of the data.id field for an element in the graph.
 * Serialized IDs are read-only and are not generated by the visualizer.
 * Cytoscape IDs are derived from Serialized IDs, though SourceSet IDs are a bit
 * more complex than that.
 */

/**
 * Highlights the selected query table row and starts the query.
 * @param {number} query_id the id of the query, corresponding to the row in the
 * query table.
*/
function select_query(query_id) {
  for (const elem of document.getElementsByClassName("datarow")) {
    elem.classList.remove("selected-row");
  }
  document.getElementById("query_row"+query_id).classList.add("selected-row");
  vis.setup_query(query_id);
}

/**
 * Visualizer handles manipulating the Cytoscape instance using the pytype
 * program data.
 * `cy` is a Cytoscape object, created by e.g. `cy = cytoscape({...});`.
 * `program` is an object created by
 * pytype.typegraph.typegraph_serializer.TypegraphEncoder, generated from a
 * pytype.typegraph.cfg.Program.
 */
class Visualizer {
  /**
   * Construct a Visualizer instance.
   * @param {!Object} cy A Cytoscape object.
   * @param {!Object} program A serialized pytype Program, generated by
   * pytype.typegraph.typegraph_serializer.TypegraphEncoder.
   */
  constructor(cy, program) {
    /** @private @const {!Object} */
    this.cy = cy;

    /** @private @const {!Object} */
    this.program = program;

    /** @private @const {!Array<!Object>} */
    this.queries = program.queries;

    /** @private @const {number?} */
    this.current_query = null;

    /** @private @const {number} */
    this.query_step = 0;

    let cfgnode_nodes = this.program.cfg_nodes.map(n => this.gen_cfgnode(n));
    let cfgnode_edges = this.program.cfg_nodes.flatMap(
        n => n.outgoing.map(
            o_id => this.gen_edge(
                'cfgnode_edge', this.cfgnode_id(n.id), this.cfgnode_id(o_id))));
    this.cy.batch(() => {
      this.cy.add(cfgnode_nodes);
      this.cy.add(cfgnode_edges);
    });
    // We immediately hide every node except the root.
    cy.nodes('[name != "root"]').toggleClass("hidden_node", true);
    this.relayout();
  }

  /**
   * Run the defined layout. This is required after adding nodes.
   */
  relayout() {
    this.cy.layout(this.cy.data('layout_options')).run();
  }

  /**
   * Set the hiddenNode style class on the Cytoscape element with the given ID.
   * @param {string} id the Cytoscape ID of the element to modify.
   */
  hide(id) {
    this.cy.$id(id).toggleClass("hidden_node", true);
  }

  /**
   * Like hide, but the opposite.
   * @param {string} id the Cytoscape ID of the element to modify.
   */
  unhide(id) {
    this.cy.$id(id).toggleClass("hidden_node", false);
  }

  /**
   * Returns the Cytoscape ID for the given CFGNode ID.
   * Note that this operates on *IDs*, not on full CFGNodes.
   * @param {number} cfgnode The ID of a SerializedCFGNode object.
   * @return {string} The Cytoscape ID for this CFGNode.
   */
  cfgnode_id(cfgnode) {
    return `cfgnode_${cfgnode}`;
  }

  /**
   * Generate a Cytoscape node for the given CFGNode.
   * @param {!Object} cfgnode A SerializedCFGNode object
   * @return {!Object} A Cytoscape element representing the CFGNode
   */
  gen_cfgnode(cfgnode) {
    return {
      group: 'nodes',
      data: {
        id: this.cfgnode_id(cfgnode.id),
        name: cfgnode.name,
        class: 'cfg_node',
      },
      classes: ['cfg_node'],
    };
  }

  /**
   * Generate the Cytoscape ID for the given Variable ID.
   * @param {number} variable The ID of a SerializedVariable object.
   * @return {string} The Cytoscape ID for the Variable.
   */
  variable_id(variable) {
    return `var_${variable}`;
  }

  /**
   * Generate a Cytoscape node for the given Variable.
   * @param {!Object} variable A SerializedVariable object.
   * @return {!Object} A Cytoscape elemenent representing the Variable.
   */
  gen_variable(variable) {
    return {
      group: 'nodes',
      data: {
        id: this.variable_id(variable.id),
        name: `v${variable.id}`,
        class: 'variable_node',
      },
      classes: ['variable_node'],
    };
  }

  /**
   * Generates the Cytoscape ID for the given Binding ID.
   * @param {number} binding The ID of a SerializedBinding object.
   * @return {string} The Cytoscape ID for the Binding.
   */
  binding_id(binding) {
    return `bind_${binding}`;
  }

  /**
   * Generate the Cytoscape node for the given Binding.
   * @param {!Object} binding A SerializedBinding object.
   * @return {!Object} A Cytoscape element representing the Binding.
   */
  gen_binding(binding) {
    return {
      group: 'nodes',
      data: {
        id: this.binding_id(binding.id),
        name: binding.data,
        class: 'binding_node',
      },
      classes: ['binding_node'],
    };
  }

  /**
   * Generate the Cytoscape element for a SourceSet.
   * @param {number} b_id Id of the Binding that this SourceSet is related to.
   * @param {number} o_id Id of the Origin that this SourceSet is part of.
   * @param {number} s_id Unique number identifying this SourceSet among the
   * other SourceSets of the Origin.
   * The idea is that, for example, the third SourceSet of the second Origin
   * of the first Binding will be generated as: gen_sourceset(0, 1, 2). This
   * is in order to create a cytoscape element with a unique ID.
   * @return {!Object} a Cytoscape element representing the SourceSet.
   */
  gen_sourceset(b_id, o_id, s_id) {
    return {
      group: 'nodes',
      data: {
        id: `sourceset_${b_id}_${o_id}_${s_id}`,
        name: '',
        class: 'sourceset_node',
      },
      classes: ['sourceset_node'],
    };
  }

  /**
   * Generic edge generator.
   * @param {string} kind The kind of edge, e.g. cfg_node_edge indicates an edge
   * between two CFGNodes, while var_bind_edge indicates an edge between a
   * Variable and a Binding. This should (but not must) match one of the edge
   * styles defined in this.cy.style.
   * @param {string} source The Cytoscape ID of the source.
   * @param {string} target The Cytoscape ID of the target.
   * @return {!Object} a Cytoscape element representing the edge.
   */
  gen_edge(kind, source, target) {
    return {
      group: 'edges',
      data: {
        id: `${kind}_${source}_${target}`,
        source: source,
        target: target,
        class: kind,
      },
      classes: [kind],
    };
  }

  /**
   * Checks if an element with the given ID already exists in the graph.
   * @param {string} elem The Cytoscape ID to check for.
   * @return {!bool} true if an element with the given ID exists.
   */
  elem_exists(elem) {
    return this.cy.$id(elem).nonempty();
  }

  /**
   * Adds all elements in the list that aren't already in the graph to the graph.
   * Cytoscape does not like it when you add an element with the same ID as an
   * existing element.
   * @param {!Array<!Object>} elems Array of Cytoscape elements to possibly add.
   */
  add_elems(elems) {
    this.cy.add(elems.filter(e => !this.elem_exists(e.data.id)));
  }

  /** Collects the node for a binding and all the edges between it and the
   * CFGNodes it is connected to.
   * @param {number} bind_id The ID of the SerializedBinding to process.
   * @return {!Array<!Object>} The list of Cytoscape elements generated for the
   * given binding.
   */
  collect_single_binding(bind_id) {
    const elems = [];
    const binding = this.program.bindings.find(b => b.id == bind_id);
    const b_node = this.gen_binding(binding);
    elems.push(b_node);
    for (const cfg_node of this.program.cfg_nodes.filter(n => n.bindings.includes(bind_id))) {
      elems.push(this.gen_edge('cfgnode_bind_edge', this.cfgnode_id(cfg_node.id), this.binding_id(bind_id)));
    }
    return elems;
  }

  /**
   * Adds a Binding and all its children to the visualization.
   * In particular, this will add:
   * - a node for the Binding.
   * - an edge between each CFGNode that has the Binding in its bindings list.
   * - a node for each SourceSet in the Binding's Origins.
   * - an edge between the Binding and each SourceSet.
   * - an edge between the SourceSet and each Binding in the SourceSet.
   * - an edge between each SourceSet and CFGNode that is connected to the
   *   Origin that contains the SourceSet.
   * It does not recurse further into the SoureceSets' member Bindings.
   * You should probably call this.relayout() after this function, before
   * calling anything else that adds nodes or edges.
   * @param {number} bind_id The ID of the SerializedBinding to process.
   */
  reveal_binding(bind_id) {
    const elems = this.collect_single_binding(bind_id);
    for (const [o_id, origin] of binding.origins.entries()) {
      for (const [s_id, sourceset] of origin.source_sets.entries()) {
        const ss = this.gen_sourceset(bind_id, o_id, s_id);
        elems.push(ss);
        elems.push(this.gen_edge('bind_source_edge', b_node.data.id, ss.data.id));
        this.unhide(this.cfgnode_id(origin.where));
        elems.push(this.gen_edge('source_cfgnode_edge', ss.data.id, this.cfgnode_id(origin.where)));
        for (const member of sourceset) {
          const mem_bind = this.program.bindings.find(b => b.id == member);
          const mem = this.gen_binding(mem_bind);
          elems.push(mem);
          elems.push(this.gen_edge('source_member_edge', ss.data.id, mem.data.id));
        }
      }
    }
    this.add_elems(elems);
  }

  /**
   * Make a Variable and all its Bindings visible in the graph.
   * In particular, this will:
   * - add a node for the Variable.
   * - add an edge from the Variable to each of its Bindings.
   * - call reveal_binding on each Binding.
   * If a Cytoscape node for the Variable has already been added, this function
   * will have no effect, unless (somehow) a new child (Binding, etc.) has
   * been added. This applies to all children of the Variable.
   * @param {number} var_id The ID number of the Variable to reveal.
   */
  reveal_var(var_id) {
    const v = this.program.variables.find(v => v.id == var_id);
    const v_node = this.gen_variable(v);
    this.add_elems([v_node]);
    const edges = [];
    for (const b_id of v.bindings) {
      this.reveal_binding(b_id);
      edges.push(this.gen_edge("var_bind_edge", v_node.data.id, this.binding_id(b_id)));
    }
    this.add_elems(edges);
    this.relayout();
  }

  /**
   * Returns the cluster of non-cfg_nodes connected to a Variable.
   * This is more useful if the Variable has been added by reveal_var already.
   * @param {number} var_id The ID number of the Variable to reveal.
   * @return {!Object} A Cytoscape.collection containing the nodes.
   */
  get_var_cluster(var_id) {
    let next = this.cy.$id(this.variable_id(var_id));
    let nodes = this.cy.collection();
    // At each step, add all the edges (and their targets) whose source is in
    // the current set of nodes.
    // We stop at the boundary of cfg_nodes to prevent this from grabbing the
    // entire graph.
    while (!nodes.same(next)) {
      nodes = next;
      // union() returns a new collection, so this is safe.
      next = next.union(next.outgoers('[class != "cfg_node"]'));
    }
    return nodes;
  }

  /**
   * Adds a Variable and its nodes to the graph, or toggles the visibility of
   * those nodes if they already exist.
   * This is intended to be used by the variable table in the visualizer.
   * @param {number} var_id The ID number of the Variable to reveal.
   */
  add_or_hide_var(var_id) {
    if (this.elem_exists(this.variable_id(var_id))) {
      this.get_var_cluster(var_id).toggleClass('hidden_node');
    } else {
      this.reveal_var(var_id);
    }
    this.get_var_cluster(var_id).flashClass('highlight_node', 1000);
  }

  /**
   * Toggles the `highlight_node` Cytoscape style class.
   * Has no effect if the node for the Variable does not exist.
   * This is intended to be used by the variable table in the visualizer.
   * There are two functions for this so it does not interfere with the
   * highlight from add_or_hide_var.
   * @param {number} var_id The ID number of the Variable to highlight.
   */
  highlight_var(var_id) {
    if (this.elem_exists(this.variable_id(var_id))) {
      this.get_var_cluster(var_id).toggleClass('highlight_node', true);
    }
  }

  /**
   * Toggles the `highlight_node` Cytoscape style class.
   * Has no effect if the node for the Variable does not exist.
   * This is intended to be used by the variable table in the visualizer.
   * @param {number} var_id The ID number of the Variable to highlight.
   */
  unhighlight_var(var_id) {
    if (this.elem_exists(this.variable_id(var_id))) {
      this.get_var_cluster(var_id).toggleClass('highlight_node', false);
    }
  }

  /**
   * Toggles on the 'highlight_node' style class for a given CFGNode.
   * Has no effect if the CFGNode id does not correspond with an actual node.
   * @param {number} node_id The ID number of the CFGNode to highlight.
   */
  highlight_cfgnode(node_id) {
    console.log("Highlighting ", this.cfgnode_id(node_id));
    this.cy.$id(this.cfgnode_id(node_id)).toggleClass('highlight_node', true);
  }

  /**
   * Toggles off the 'highlight_node' style class for a given CFGNode.
   * Has no effect if the CFGNode id does not correspond with an actual node.
   * @param {number} node_id The ID number of the CFGNode to highlight.
   */
  unhighlight_cfgnode(node_id) {
    this.cy.$id(this.cfgnode_id(node_id)).toggleClass('highlight_node', false);
  }

  /** Sets up a query for visualization.
   * - Unsets the current query, if there is one.
   * - Visualizes the first step of the query.
   * @param {number} query_idx The index of the query to setup.
  */
  setup_query(query_idx) {
    if (this.queries[query_idx] === undefined) {
      console.log("Query", query_idx, "is not a valid query index:", this.queries);
      return;
    }

    if (this.current_query === query_idx) {
      return;
    }

    const new_query = this.queries[query_idx];

    // Reset the current query, if there is one.
    if (this.current_query !== null) {
      const query = this.queries[this.current_query];
      for (const step of query.steps.slice(0, this.query_step+1)) {
        this.unhighlight_cfgnode(step.node);
        // Hide all the nodes for the query, unless it's visited by the new
        // query. This avoids a bug related to toggling the same class multiple
        // times between relayout() calls.
        if (new_query.steps.find(s => s.node !== step.node)) {
          this.hide(this.cfgnode_id(step.node));
        }
        for (const b_id of step.bindings) {
            this.hide(this.binding_id(b_id));
        }
      }
    }

    this.current_query = query_idx;
    for (const step of new_query.steps) {
      this.unhide(this.cfgnode_id(step.node));
    }

    // query_step is immediately incremented by advance_query.
    this.query_step = -1;
    this.advance_query();
  }

  /**
   * Show the next step of a query.
   * Does nothing if the query is already at the final step.
  */
  advance_query() {
    if (this.current_query === null) {
      return;
    }

    if (this.query_step === this.queries[this.current_query].steps.length-1) {
      return;
    }

    // Hide the previous step, if there is one.
    if (this.query_step >= 0) {
      const step = this.queries[this.current_query].steps[this.query_step];
      this.unhighlight_cfgnode(step.node);
      for (const b_id of step.bindings) {
          this.hide(this.binding_id(b_id));
      }
    }

    this.query_step += 1;
    const step = this.queries[this.current_query].steps[this.query_step];
    this.highlight_cfgnode(step.node);
    const elems = [];
    // We don't know if a binding has been revealed before, so we need to both
    // reveal the bindings *and* un-hide them.
    for (const bind_id of step.bindings) {
      elems.push(...this.collect_single_binding(bind_id));
    }
    this.add_elems(elems);
    for (const elem of elems) {
      this.unhide(elem.data.id);
    }
    this.relayout();
  }

  /**
   * Show the previous step of a query.
   * Does nothing if the query is already at the first step.
  */
  retreat_query() {
    if (this.current_query === null) {
      return;
    }

    if (this.query_step === 0) {
      return;
    }

    
    let curr_step = this.queries[this.current_query].steps[this.query_step];
    this.query_step -= 1;
    let new_step = this.queries[this.current_query].steps[this.query_step];

    // first, unhighlight/hide the current step. Cytoscape (or dagre) doesn't
    // like it if we toggle a class without calling relayout() in between, so
    // only hide a binding if it isn't also present in the step we're
    // transitioning to.
    this.unhighlight_cfgnode(curr_step.node);
    for (const b_id of curr_step.bindings) {
      if (!new_step.bindings.includes(b_id)) {
        this.hide(this.binding_id(b_id));
      }
    }

    // then process the previous step. All the nodes have already been added
    // to the graph, so they just need to be unhidden.
    this.highlight_cfgnode(new_step.node);
    for (const bind_id of new_step.bindings) {
      for (const elem of this.collect_single_binding(bind_id)) {
        this.unhide(elem.data.id);
      }
    }
    this.relayout();
  }

  /**
   * Generates the Cytoscape elements for the legend.
   * These should be used in another Cytoscape instance.
   * @return {!Array<!Object>} List of elements usable by Cytoscape.
   */
  gen_legend() {
    // Need to set manually set names on SourceSets.
    const num_sourcesets = 3;
    const sourcesets = [];
    for (let index = 0; index < num_sourcesets; index++) {
      let ss = this.gen_sourceset(0, 0, index);
      ss.data.name = 'SourceSet';
      sourcesets.push(ss);
    }

    // Same with variables, but we only have the one.
    const variable = this.gen_variable({id: 0});
    variable.data.name = 'Variable';

    return [
      this.gen_cfgnode({id: 0, name: 'CFGNode'}),
      this.gen_cfgnode({id: 1, name: 'CFGNode'}),
      this.gen_edge('cfgnode_edge', this.cfgnode_id(0), this.cfgnode_id(1)),

      variable,
      this.gen_binding({id: 0, data: 'Binding'}),
      this.gen_edge('var_bind_edge', this.variable_id(0), this.binding_id(0)),

      sourcesets[0],
      this.gen_cfgnode({id: 2, name: 'CFGNode'}),
      this.gen_edge('source_cfgnode_edge', sourcesets[0].data.id, this.cfgnode_id(2)),

      this.gen_binding({id: 1, data: 'Binding'}),
      sourcesets[1],
      this.gen_edge('bind_source_edge', this.binding_id(1), sourcesets[1].data.id),

      sourcesets[2],
      this.gen_binding({id: 2, data: 'Member'}),
      this.gen_edge('source_member_edge', sourcesets[2].data.id, this.binding_id(2)),

      this.gen_cfgnode({id: 3, name: 'CFGNode'}),
      this.gen_binding({id: 3, data: 'Binding'}),
      this.gen_edge('cfgnode_bind_edge', this.cfgnode_id(3), this.binding_id(3)),
    ];
  }
}